import colorsys
import os
import time

import numpy as np
import torch
import torch.backends.cudnn as cudnn
from PIL import ImageDraw, ImageFont

from nets.centernet import CenterNet_HourglassNet, CenterNet_Resnet50
from utils.utils import cvtColor, get_classes, preprocess_input, resize_image
from utils.utils_bbox import decode_bbox, postprocess,draw_heatmap
from pyheatmap.heatmap import HeatMap
from PIL import Image
import cv2
import matplotlib.pyplot as plt

def apply_heatmap(image,data,r=5):
    #image是原图，data是坐标
    #创建一个新的与原图大小一致的图像，color为0背景为黑色。这里这样做是因为在绘制热力图的时候如果不选择背景图，画出来的图与原图大小不一致（根据点的坐标来的），导致无法对热力图和原图进行加权叠加，因此，这里我新建了一张背景图。
    background = Image.new("RGB", (image.shape[1], image.shape[0]),color=(125,125,125))
    # 开始绘制热度图
    hm = HeatMap(data)
    hit_img = hm.heatmap(base=background, r = r) # background为背景图片，r是半径，默认为10
    #plt.figure()
    #plt.imshow(hit_img)
    #plt.show()
    #hit_img.save('out_' + image_name + '.jpeg')
    hit_img = cv2.cvtColor(np.asarray(hit_img),cv2.COLOR_RGB2BGR)#Image格式转换成cv2格式
    hit_img = cv2.cvtColor(hit_img, cv2.COLOR_BGR2RGB)
    overlay = image.copy()
    alpha = 0.5 # 设置覆盖图片的透明度
    cv2.rectangle(overlay, (0, 0), (image.shape[1], image.shape[0]), (0, 0, 255), -1) # 设置蓝色为热度图基本色蓝色
    image = cv2.addWeighted(overlay, alpha, image, 1-alpha, 0) # 将背景热度图覆盖到原图
    try:
        image = cv2.addWeighted(hit_img, 1 - alpha, image, alpha, 0)  # 将热度图覆盖到原图
    except:
        print()
    return image
#--------------------------------------------#
#   使用自己训练好的模型预测需要修改3个参数
#   model_path、classes_path和backbone
#   都需要修改！
#   如果出现shape不匹配，一定要注意
#   训练时的model_path和classes_path参数的修改
#--------------------------------------------#
class CenterNet(object):
    _defaults = {
        #--------------------------------------------------------------------------#
        #   使用自己训练好的模型进行预测一定要修改model_path和classes_path！
        #   model_path指向logs文件夹下的权值文件，classes_path指向model_data下的txt
        #
        #   训练好后logs文件夹下存在多个权值文件，选择验证集损失较低的即可。
        #   验证集损失较低不代表mAP较高，仅代表该权值在验证集上泛化性能较好。
        #   如果出现shape不匹配，同时要注意训练时的model_path和classes_path参数的修改
        #--------------------------------------------------------------------------#

        "model_path"        : 'model_data/mban_model_weight.pth',

        "classes_path"      : 'model_data/people_classes.txt',
        #--------------------------------------------------------------------------#
        #   用于选择所使用的模型的主干
        #   resnet50, hourglass
        #--------------------------------------------------------------------------#
        "backbone"          : 'resnet50',
        #--------------------------------------------------------------------------#
        #   输入图片的大小，设置成32的倍数
        #--------------------------------------------------------------------------#
        "input_shape"       : [416, 416],
        #--------------------------------------------------------------------------#
        #   只有得分大于置信度的预测框会被保留下来
        #--------------------------------------------------------------------------#
        "confidence"        : 0.5,
        #---------------------------------------------------------------------#
        #   非极大抑制所用到的nms_iou大小
        #---------------------------------------------------------------------#
        "nms_iou"           : 0.3,
        #--------------------------------------------------------------------------#
        #   是否进行非极大抑制，可以根据检测效果自行选择
        #   backbone为resnet50时建议设置为True、backbone为hourglass时建议设置为False
        #--------------------------------------------------------------------------#
        "nms"               : True,
        #---------------------------------------------------------------------#
        #   该变量用于控制是否使用letterbox_image对输入图像进行不失真的resize，
        #   在多次测试后，发现关闭letterbox_image直接resize的效果更好
        #---------------------------------------------------------------------#
        "letterbox_image"   : False,
        #-------------------------------#
        #   是否使用Cuda
        #   没有GPU可以设置成False
        #-------------------------------#
        "cuda"              : True
    }

    @classmethod
    def get_defaults(cls, n):
        if n in cls._defaults:
            return cls._defaults[n]
        else:
            return "Unrecognized attribute name '" + n + "'"

    #---------------------------------------------------#
    #   初始化centernet
    #---------------------------------------------------#
    def __init__(self, **kwargs):
        self.__dict__.update(self._defaults)
        for name, value in kwargs.items():
            setattr(self, name, value)
        #---------------------------------------------------#
        #   计算总的类的数量
        #---------------------------------------------------#
        self.class_names, self.num_classes  = get_classes(self.classes_path)

        #---------------------------------------------------#
        #   画框设置不同的颜色
        #---------------------------------------------------#
        hsv_tuples = [(x / self.num_classes, 1., 1.) for x in range(self.num_classes)]
        self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
        self.colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), self.colors))

        self.generate()

    #---------------------------------------------------#
    #   载入模型
    #---------------------------------------------------#
    def generate(self):
        #-------------------------------#
        #   载入模型与权值
        #-------------------------------#
        assert self.backbone in ['resnet50', 'hourglass']
        if self.backbone == "resnet50":
            self.net = CenterNet_Resnet50(num_classes=self.num_classes, pretrained=False)
        else:
            self.net = CenterNet_HourglassNet({'hm': self.num_classes, 'wh': 2, 'reg':2})

        device      = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.net.load_state_dict(torch.load(self.model_path, map_location=device))
        self.net    = self.net.eval()
        print('{} model, and classes loaded.'.format(self.model_path))

        if self.cuda:
            self.net = torch.nn.DataParallel(self.net)
            cudnn.benchmark = True
            self.net = self.net.cuda()

    def heat_image(self,image):
        # ---------------------------------------------------#
        #   计算输入图片的高和宽
        # ---------------------------------------------------#
        image_shape = np.array(np.shape(image)[0:2])

        # ---------------------------------------------------------#
        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
        # ---------------------------------------------------------#
        image = cvtColor(image)
        # ---------------------------------------------------------#
        #   给图像增加灰条，实现不失真的resize
        #   也可以直接resize进行识别
        # ---------------------------------------------------------#
        image_data = resize_image(image, (self.input_shape[1], self.input_shape[0]), self.letterbox_image)

        # -----------------------------------------------------------#
        #   图片预处理，归一化。获得的photo的shape为[1, 512, 512, 3]
        # -----------------------------------------------------------#
        image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)

        with torch.no_grad():
            images = torch.from_numpy(np.asarray(image_data)).type(torch.FloatTensor)
            if self.cuda:
                images = images.cuda()
            # ---------------------------------------------------------#
            #   将图像输入网络当中进行预测！
            #   hm1   =[1,numclass,104,104]
            #   hm2   =[1,numclass,104,104]
            #   wh    =[1,2,104,104]
            #   offset=[1,2,104,104]
            # ---------------------------------------------------------#
            outputs = self.net(images)

            if self.backbone == 'hourglass':
                outputs = [outputs[-1]["hm"].sigmoid(), outputs[-1]["wh"], outputs[-1]["reg"]]
            # -----------------------------------------------------------#
            #   利用预测结果进行解码
            #   output1=(1)
            # -----------------------------------------------------------#
            outputs1 = decode_bbox(outputs[0],outputs[1], outputs[2], outputs[3],outputs[4], self.confidence, self.cuda)

            # outputs2 = decode_bbox(outputs[1], outputs[2], outputs[3], self.confidence, self.cuda)
            # -------------------------------------------------------#
            #   对于centernet网络来讲，确立中心非常重要。
            #   对于大目标而言，会存在许多的局部信息。
            #   此时对于同一个大目标，中心点比较难以确定。
            #   使用最大池化的非极大抑制方法无法去除局部框
            #   所以我还是写了另外一段对框进行非极大抑制的代码
            #   实际测试中，hourglass为主干网络时有无额外的nms相差不大，resnet相差较大。
            # -------------------------------------------------------#

            results = postprocess(outputs1, self.nms, image_shape, self.input_shape, self.letterbox_image, self.nms_iou)

            # --------------------------------------#
            #   如果没有检测到物体，则返回原图
            # --------------------------------------#
            if results[0] is None:
                return image

            top_label = np.array(results[0][:, 5], dtype='int32')
            top_conf = results[0][:, 4]
            top_boxes = results[0][:, :4]
            center_points = (top_boxes[:, 0:2] + top_boxes[:, 2:4]) / 2

            center_points=center_points.astype(np.int)
            print(center_points)
            imge=draw_heatmap(image,center_points)
            plt.imshow(imge)
            plt.xticks([])  # 去掉横坐标值
            plt.yticks([])  # 去掉纵坐标值
            #plt.savefig('img/%03d-%03d.png' % (epoch + 1, iteration + 1))
            plt.show()
            plt.close()

    #---------------------------------------------------#
    #   检测图片
    #---------------------------------------------------#
    def detect_image(self, image):
        #---------------------------------------------------#
        #   计算输入图片的高和宽
        #---------------------------------------------------#
        image_shape = np.array(np.shape(image)[0:2])

        #---------------------------------------------------------#
        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
        #---------------------------------------------------------#
        image       = cvtColor(image)
        #---------------------------------------------------------#
        #   给图像增加灰条，实现不失真的resize
        #   也可以直接resize进行识别
        #---------------------------------------------------------#
        image_data  = resize_image(image, (self.input_shape[1], self.input_shape[0]), self.letterbox_image)

        #-----------------------------------------------------------#
        #   图片预处理，归一化。获得的photo的shape为[1, 512, 512, 3]
        #-----------------------------------------------------------#
        image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)

        with torch.no_grad():
            images = torch.from_numpy(np.asarray(image_data)).type(torch.FloatTensor)
            if self.cuda:
                images = images.cuda()
            #---------------------------------------------------------#
            #   将图像输入网络当中进行预测！
            #   hm1   =[1,numclass,104,104]
            #   hm2   =[1,numclass,104,104]
            #   hm3   =[1,numclass,104,104]
            #   wh    =[1,2,104,104]
            #   offset=[1,2,104,104]
            #---------------------------------------------------------#
            outputs = self.net(images)

            if self.backbone == 'hourglass':
                outputs = [outputs[-1]["hm"].sigmoid(), outputs[-1]["wh"], outputs[-1]["reg"]]
            #-----------------------------------------------------------#
            #   利用预测结果进行解码
            #   output1=(1)
            #-----------------------------------------------------------#
            outputs,outputs1,outputs2 = decode_bbox(outputs[0],outputs[1], outputs[2], outputs[3],outputs[4], self.confidence, self.cuda)

            #outputs2 = decode_bbox(outputs[1], outputs[2], outputs[3], self.confidence, self.cuda)
            #-------------------------------------------------------#
            #   对于centernet网络来讲，确立中心非常重要。
            #   对于大目标而言，会存在许多的局部信息。
            #   此时对于同一个大目标，中心点比较难以确定。
            #   使用最大池化的非极大抑制方法无法去除局部框
            #   所以我还是写了另外一段对框进行非极大抑制的代码
            #   实际测试中，hourglass为主干网络时有无额外的nms相差不大，resnet相差较大。
            #-------------------------------------------------------#
            resultss = postprocess(outputs, self.nms, image_shape, self.input_shape, self.letterbox_image,                 self.nms_iou)
            resultss1 = postprocess(outputs1, self.nms, image_shape, self.input_shape, self.letterbox_image, self.nms_iou)
            resultss2 = postprocess(outputs2, self.nms, image_shape, self.input_shape, self.letterbox_image, self.nms_iou)
            #--------------------------------------#
            #   如果没有检测到物体，则返回原图
            #--------------------------------------#
            if resultss is None  :
                return image
            if resultss1 is None  :
                return image
            if resultss2 is None  :
                return image
            # try:
            #     results_midbox = resultss[0]
            #     results_upbox = resultss1[0]
            #     results_doownbox = resultss2[0]
            #     results_midbox = results_midbox.tolist()
            #     results_upbox = results_upbox.tolist()
            #     results_downbox = results_doownbox.tolist()
            #     # print(np.shape(results_midbox))
            #     # print(np.shape(results_downbox))
            #     # print(np.shape(results_upbox))
            #     count_up = []
            #     count_down = []
            #     for m in range(len(results_midbox)):
            #         mid_box=results_midbox[m]
            #         mid_y = (mid_box[0] + mid_box[2]) / 2
            #         mid_x = (mid_box[1] + mid_box[3]) / 2
            #         for n in range(len(results_upbox)):
            #             up_box = results_upbox[n]
            #
            #             up_y = (up_box[0] + up_box[2]) / 2
            #             up_x = (up_box[1] + up_box[3]) / 2
            #
            #             if up_box[1] > mid_box[1]  and up_box[2] < mid_box[2]  and up_box[3] <mid_box[3]:
            #                 if np.sqrt(np.square(up_x-mid_x)+np.square(up_y-mid_y))<((mid_box[0]+mid_box[2])/2):
            #                     if ((up_box[3] - up_box[1]) * (up_box[2] - mid_box[0])) / ((up_box[2] - up_box[0]) * (up_box[3] - up_box[1])) > 0.7:
            #                         count_up.append(n)
            #                         #break
            #
            #         for k in range(len(results_downbox)):
            #             down_box = results_downbox[k]
            #             down_y = (down_box[0] + down_box[2]) / 2
            #             down_x = (down_box[1] + down_box[3]) / 2
            #             if down_box[0] > mid_box[0] and down_box[1] > mid_box[1] and down_box[2] < mid_box[2] and down_box[3] < \
            #                     mid_box[3]:
            #                 if mid_x * 0.65 < down_x < mid_x * 1.45:
            #                     count_down.append(k)
            #                     break
            #     results_downboxx = [results_downbox[i] for i in range(0, len(results_downbox), 1) if i not in count_down]
            #     results_upboxx = [results_upbox[i] for i in range(0, len(results_upbox), 1) if i not in count_up]
            #
            #     new_count=[]
            #     for i in range(len(results_downboxx)):
            #         down_box = results_downboxx[i]
            #         down_y = (down_box[0] + down_box[2]) / 2
            #         down_x = (down_box[1] + down_box[3]) / 2
            #         min_dis = 9999
            #         min_dis1 = 9999
            #         temp=0
            #         temp1 = 0
            #         for j in range(len(results_upboxx)):
            #             # if j in count:
            #             #     continue
            #             up_box = results_upboxx[j]
            #             up_y = (up_box[0] + up_box[2]) / 2
            #             up_x = (up_box[1] + up_box[3]) / 2
            #
            #             if  up_y<down_y :
            #                 sta = np.abs(np.arctan((down_y-up_y)/((down_x-up_x)+ 1e-4))*(180 / np.pi))
            #                 if sta<30:
            #                     dist=np.sqrt(np.square(up_x-down_x)+np.square(up_y-down_y))
            #                     if dist<min_dis:
            #                         min_dis=dist
            #                         temp=j
            #             if  up_x<down_x :
            #                 sta = np.abs(np.arctan((down_x-up_x)/((down_y-up_y)+ 1e-4))*(180 / np.pi))
            #                 if sta<30:
            #                     dist=np.sqrt(np.square(up_y-down_y)+np.square(up_x-down_x))
            #                     if dist<min_dis1:
            #                         min_dis=dist
            #                         temp1=j
            #         new_count.append(temp)
            #         new_count.append(temp1)
            #
            #
            #     # results_box = results_box.tolist()
            #     results_downboxxx = [results_downboxx[i] for i in range(0, len(results_downboxx), 1) if i not in new_count]
            #     #results_upboxxx = [results_upboxx[i] for i in range(0, len(results_upboxx), 1) if i not in count_up]
            #     results_box=results_midbox+results_upboxx+results_downboxxx
            #     results_box = np.array(results_box)
            # except:
            #     return image
            results_midbox = resultss[0]
            results_upbox = resultss1[0]
            results_doownbox = resultss2[0]
            results_midbox = results_midbox.tolist()
            results_upbox = results_upbox.tolist()
            results_downbox = results_doownbox.tolist()
            results_box = results_midbox + results_upbox  + results_downbox
            results_box = np.array(results_box)
            top_label = np.array(results_box[:, 5], dtype='int32')
            top_conf = results_box[:, 4]
            top_boxes = results_box[:, :4]

            #center_points=(top_boxes[:,0:2]+top_boxes[:,2:4])/2


        #---------------------------------------------------------#
        #   设置字体与边框厚度
        #---------------------------------------------------------#
        font = ImageFont.truetype(font='model_data/simhei.ttf', size=np.floor(3e-2 * np.shape(image)[1] + 0.5).astype('int32'))
        thickness = max((np.shape(image)[0] + np.shape(image)[1]) // self.input_shape[0], 1)

        #---------------------------------------------------------#
        #   图像绘制
        #---------------------------------------------------------#
        for i, c in list(enumerate(top_label)):
            predicted_class = self.class_names[int(c)]
            box             = top_boxes[i]
            score           = top_conf[i]

            top, left, bottom, right = box

            top     = max(0, np.floor(top).astype('int32'))
            left    = max(0, np.floor(left).astype('int32'))
            bottom  = min(image.size[1], np.floor(bottom).astype('int32'))
            right   = min(image.size[0], np.floor(right).astype('int32'))

            label = '{} {:.2f}'.format(predicted_class, score)
            draw = ImageDraw.Draw(image)
            label_size = draw.textsize(label, font)
            label = label.encode('utf-8')
            print(label, top, left, bottom, right)
            
            if top - label_size[1] >= 0:
                text_origin = np.array([left, top - label_size[1]])
            else:
                text_origin = np.array([left, top + 1])

            for i in range(thickness):
                draw.rectangle([left + i, top + i, right - i, bottom - i], outline=self.colors[c])
            #draw.rectangle([tuple(text_origin), tuple(text_origin + label_size)], fill=self.colors[c])
            #draw.text(text_origin, str(label,'UTF-8'), fill=(0, 0, 0), font=font)
            del draw
            # print(box)
            # plt.imshow(image)
            # plt.show()
            # plt.close()
        return image


    def get_FPS(self, image, test_interval):
        #---------------------------------------------------#
        #   计算输入图片的高和宽
        #---------------------------------------------------#
        image_shape = np.array(np.shape(image)[0:2])
        #---------------------------------------------------------#
        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
        #---------------------------------------------------------#
        image       = cvtColor(image)
        #---------------------------------------------------------#
        #   给图像增加灰条，实现不失真的resize
        #   也可以直接resize进行识别
        #---------------------------------------------------------#
        image_data  = resize_image(image, (self.input_shape[1], self.input_shape[0]), self.letterbox_image)
        #-----------------------------------------------------------#
        #   图片预处理，归一化。获得的photo的shape为[1, 512, 512, 3]
        #-----------------------------------------------------------#
        image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)

        with torch.no_grad():
            images = torch.from_numpy(np.asarray(image_data)).type(torch.FloatTensor)
            if self.cuda:
                images = images.cuda()
            #---------------------------------------------------------#
            #   将图像输入网络当中进行预测！
            #---------------------------------------------------------#
            outputs = self.net(images)
            if self.backbone == 'hourglass':
                outputs = [outputs[-1]["hm"].sigmoid(), outputs[-1]["wh"], outputs[-1]["reg"]]
            #-----------------------------------------------------------#
            #   利用预测结果进行解码
            #-----------------------------------------------------------#
            #outputs = decode_bbox(outputs[0], outputs[1], outputs[2], self.confidence, self.cuda)
            outputs=    decode_bbox(outputs[0], outputs[1], outputs[2], outputs[3], outputs[4], self.confidence, self.cuda)
            #-------------------------------------------------------#
            #   对于centernet网络来讲，确立中心非常重要。
            #   对于大目标而言，会存在许多的局部信息。
            #   此时对于同一个大目标，中心点比较难以确定。
            #   使用最大池化的非极大抑制方法无法去除局部框
            #   所以我还是写了另外一段对框进行非极大抑制的代码
            #   实际测试中，hourglass为主干网络时有无额外的nms相差不大，resnet相差较大。
            #-------------------------------------------------------#
            results = postprocess(outputs, self.nms, image_shape, self.input_shape, self.letterbox_image, self.nms_iou)
            
        #t1 = time.time()
        for i in range(test_interval):
            if i==5:
                t1 = time.time()
            with torch.no_grad():
                #---------------------------------------------------------#
                #   将图像输入网络当中进行预测！
                #---------------------------------------------------------#
                outputs = self.net(images)
                # if self.backbone == 'hourglass':
                #     outputs = [outputs[-1]["hm"].sigmoid(), outputs[-1]["wh"], outputs[-1]["reg"]]
                #-----------------------------------------------------------#
                #   利用预测结果进行解码
                #-----------------------------------------------------------#
                #outputs = decode_bbox(outputs[0], outputs[1], outputs[2], self.confidence, self.cuda)
                # outputs = decode_bbox(outputs[0], outputs[1], outputs[2], outputs[3], outputs[4], self.confidence,
                #                       self.cuda)
                #-------------------------------------------------------#
                #   对于centernet网络来讲，确立中心非常重要。
                #   对于大目标而言，会存在许多的局部信息。
                #   此时对于同一个大目标，中心点比较难以确定。
                #   使用最大池化的非极大抑制方法无法去除局部框
                #   所以我还是写了另外一段对框进行非极大抑制的代码
                #   实际测试中，hourglass为主干网络时有无额外的nms相差不大，resnet相差较大。
                #-------------------------------------------------------#
                #results = postprocess(outputs, self.nms, image_shape, self.input_shape, self.letterbox_image, self.nms_iou)
        t2 = time.time()
        tact_time = (t2 - t1) / (test_interval-5)
        return tact_time
        
    def get_map_txt(self, image_id, image, class_names, map_out_path):
        f = open(os.path.join(map_out_path, "detection-results/"+image_id+".txt"),"w") 
        #---------------------------------------------------#
        #   计算输入图片的高和宽
        #---------------------------------------------------#
        image_shape = np.array(np.shape(image)[0:2])
        #---------------------------------------------------------#
        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
        #---------------------------------------------------------#
        image       = cvtColor(image)
        #---------------------------------------------------------#
        #   给图像增加灰条，实现不失真的resize
        #   也可以直接resize进行识别
        #---------------------------------------------------------#
        image_data  = resize_image(image, (self.input_shape[1], self.input_shape[0]), self.letterbox_image)
        #-----------------------------------------------------------#
        #   图片预处理，归一化。获得的photo的shape为[1, 512, 512, 3]
        #-----------------------------------------------------------#
        image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)

        with torch.no_grad():
            images = torch.from_numpy(np.asarray(image_data)).type(torch.FloatTensor)
            if self.cuda:
                images = images.cuda()
            #---------------------------------------------------------#
            #   将图像输入网络当中进行预测！
            #---------------------------------------------------------#
            outputs = self.net(images)
            if self.backbone == 'hourglass':
                outputs = [outputs[-1]["hm"].sigmoid(), outputs[-1]["wh"], outputs[-1]["reg"]]
            #-----------------------------------------------------------#
            #   利用预测结果进行解码
            #-----------------------------------------------------------#
            #outputs = decode_bbox(outputs[0], outputs[1], outputs[2], self.confidence, self.cuda)
            outputs = decode_bbox(outputs[0], outputs[1], outputs[2], outputs[3], outputs[4], self.confidence,
                                  self.cuda)
            #-------------------------------------------------------#
            #   对于centernet网络来讲，确立中心非常重要。
            #   对于大目标而言，会存在许多的局部信息。
            #   此时对于同一个大目标，中心点比较难以确定。
            #   使用最大池化的非极大抑制方法无法去除局部框
            #   所以我还是写了另外一段对框进行非极大抑制的代码
            #   实际测试中，hourglass为主干网络时有无额外的nms相差不大，resnet相差较大。
            #-------------------------------------------------------#
            results = postprocess(outputs, self.nms, image_shape, self.input_shape, self.letterbox_image, self.nms_iou)
            
            #--------------------------------------#
            #   如果没有检测到物体，则返回原图
            #--------------------------------------#
            if results[0] is None:
                return 

            top_label   = np.array(results[0][:, 5], dtype = 'int32')
            top_conf    = results[0][:, 4]
            top_boxes   = results[0][:, :4]


        for i, c in list(enumerate(top_label)):
            predicted_class = self.class_names[int(c)]
            box             = top_boxes[i]
            score           = str(top_conf[i])
            
            top, left, bottom, right = box

            if predicted_class not in class_names:
                continue

            f.write("%s %s %s %s %s %s\n" % (predicted_class, score[:6], str(int(left)), str(int(top)), str(int(right)),str(int(bottom))))

        f.close()
        return 
